import tensorflow as tf
import numpy as np
from math import sin, cos, radians, sqrt

# 初始化参数
bit_num = 1
# nRow = 32  # x方向单元个数，θ方向
# mLine = 32  # y方向单元个数，φ方向
nRow = 48  # x方向单元个数，θ方向
mLine = 48  # y方向单元个数，φ方向
# nRow = 64  # x方向单元个数，θ方向
# mLine = 64  # y方向单元个数，φ方向
f0 = 300  # 工作频点，单位：GHz
c = 3e8  # 光速，单位：m/s
lamda0 = c / (f0 * 1e9)  # 波长，单位：m
dx = lamda0 / 2  # x方向单元间距，单位：m
dy = dx  # y方向单元间距，单位：m
k = 2 * np.pi / lamda0  # 波数

# 配置馈元和计算馈元位置
InF = 1  # 喇叭馈源摆放状态，0：垂直阵面，1：指向阵面中心
InTetha = 0  # 入射角
aphea_10dB = 45  # 喇叭馈源的-10dB波束角宽度

# 计算辐射方向图
qf = 1
qe = 1
eps = 0.000001


# 计算距离
def distance(p1, p2):
    return tf.sqrt(tf.reduce_sum((p1 - p2) ** 2))


# 计算馈元
def get_feed(InF, InTetha, aphea_10dB, nRow, dx):
    D = nRow * dx
    if InF == 0:
        beata = 90 - aphea_10dB / 2
        gema = 90 - beata - InTetha
        F = sin(radians(beata)) * ((D / 2) / sin(radians(gema)))
        xf = F * sin(radians(InTetha))
        yf = 0
        zf = F * cos(radians(InTetha))
    elif InF == 1:
        beata = 90 - InTetha
        gema = 180 - beata - aphea_10dB / 2
        F = sin(radians(gema)) * ((D / 2) / sin(radians(aphea_10dB / 2)))
        xf = F * sin(radians(InTetha))
        yf = 0
        zf = F * cos(radians(InTetha))
    return [xf, yf, zf], F / D


# 计算方向图, fft方式, xyz坐标
def get_pattern_fft(phase, posx, posy, k, Feed, Rc, qf, qe):
    Nx, Ny = len(posx), len(posy)
    Kfft, Lfft = 360, 361

    # 生成相位矩阵，这里假设已经得到了正确的相位值
    phase_matrix = tf.exp(tf.complex(real=0.0, imag=phase))

    # 根据波长和阵元位置计算u和v
    u = tf.linspace(-1., 1., Kfft)
    v = tf.linspace(-1., 1., Lfft)
    u, v = tf.meshgrid(u, v, indexing='ij')

    # 计算每个单元到馈源的距离差分项
    distance_diff = tf.zeros((Nx, Ny), dtype=tf.float32)
    for i in range(Nx):
        for j in range(Ny):
            E = tf.stack([posx[i], posy[j], tf.constant(0., dtype=tf.float32)])
            rfij = distance(Feed, E)
            distance_diff = tf.tensor_scatter_nd_update(distance_diff, [[i, j]], [-k * rfij])

    # 将距离差分项加到相位矩阵中
    phase_matrix *= tf.exp(tf.complex(real=0.0, imag=distance_diff))

    # 使用FFT进行快速计算
    IF = tf.signal.fftshift(tf.signal.ifft2d(phase_matrix))
    F = IF * Kfft * Lfft  # 这里乘以采样点数是为了匹配直接计算的结果尺度

    # 返回绝对值作为方向图
    return tf.abs(F), u, v  # 注意这里的u, v代替了原来的th, ph


# 相位生成方向图, xy坐标系, FFT方式
def phase_2_pattern_xyz_fft(phaseBit):
    # 1. 相位生成方向图, 极坐标
    # 生成天线位置坐标
    posx = tf.range(-dx * (nRow - 1) / 2, dx * (nRow - 1) / 2 + dx, delta=dx)
    posy = tf.range(-dy * (mLine - 1) / 2, dy * (mLine - 1) / 2 + dy, delta=dy)
    # 参考阵元位置
    Rc = tf.stack([posx[0], posy[0], tf.constant(0., dtype=tf.float32)])
    # 计算馈元位置
    Feed, FD = get_feed(InF, InTetha, aphea_10dB, nRow, dx)
    Feed = tf.constant(Feed, dtype=tf.float32)
    # 计算方向图
    pattern, th, ph = get_pattern_fft(phaseBit, posx, posy, k, Feed, Rc, qf, qe)
    # 2. 转换为直角坐标系
    x = th
    y = ph
    return pattern, x, y


# 计算phase的方向图
def phaseBit_2_pattern(phaseBit):
    # 相位转角度
    phaseBitDeg = phaseBit * 180.
    # 角度转弧度
    pi = tf.constant(np.pi, dtype=tf.float32)
    phaseBitRad = phaseBitDeg * (pi / 180.)
    pattern_xyz, x, y = phase_2_pattern_xyz_fft(phaseBitRad)  # FFT法计算, 快速
    return pattern_xyz, x, y


@tf.function
def pattern_loss(y_true, y_pred):
    # 首先去除多余的维度 (batchsize, 64, 64, 1) -> (batchsize, 64, 64)
    y_true = tf.squeeze(y_true, axis=-1)
    y_pred = tf.squeeze(y_pred, axis=-1)
    # 获取batch size
    batch_size = tf.shape(y_true)[0]
    # 初始化总损失
    total_loss = tf.constant(0.0, dtype=tf.float32)
    # 对batch中的每个样本单独计算损失
    for i in tf.range(batch_size):
        # 获取当前样本
        current_true = y_true[i]
        current_pred = y_pred[i]
        # 计算方向图
        true_pattern, _, _ = phaseBit_2_pattern(current_true)
        pred_pattern, _, _ = phaseBit_2_pattern(current_pred)
        # 累加损失
        total_loss += tf.reduce_sum(tf.abs(true_pattern - pred_pattern))
    # 计算平均损失
    avg_loss = total_loss / tf.cast(batch_size, tf.float32)
    return avg_loss


@tf.function
def pattern_loss_correlation(y_true, y_pred):
    # 首先去除多余的维度 (batchsize, 64, 64, 1) -> (batchsize, 64, 64)
    y_true = tf.squeeze(y_true, axis=-1)
    y_pred = tf.squeeze(y_pred, axis=-1)
    # 获取batch size
    batch_size = tf.shape(y_true)[0]
    # 初始化总损失
    total_loss = tf.constant(0.0, dtype=tf.float32)
    # 对batch中的每个样本单独计算损失
    for i in tf.range(batch_size):
        # 获取当前样本
        current_true = y_true[i]
        current_pred = y_pred[i]
        # 计算方向图
        true_pattern, _, _ = phaseBit_2_pattern(current_true)
        pred_pattern, _, _ = phaseBit_2_pattern(current_pred)
        # 归一化
        true_norm = (true_pattern - tf.reduce_mean(true_pattern)) / (tf.math.reduce_std(true_pattern) + 1e-8)
        pred_norm = (pred_pattern - tf.reduce_mean(pred_pattern)) / (tf.math.reduce_std(pred_pattern) + 1e-8)
        # 相关系数
        corr = tf.reduce_sum(true_norm * pred_norm) / tf.sqrt(
            tf.reduce_sum(true_norm ** 2) * tf.reduce_sum(pred_norm ** 2)
        )
        # 转换为损失（越小越好）
        loss = 1.0 - corr
        total_loss += loss
    # 计算平均损失
    avg_loss = total_loss / tf.cast(batch_size, tf.float32)
    return avg_loss


@tf.function
def pattern_loss_ssim(y_true, y_pred):
    # 首先去除多余的维度 (batchsize, 64, 64, 1) -> (batchsize, 64, 64)
    y_true = tf.squeeze(y_true, axis=-1)
    y_pred = tf.squeeze(y_pred, axis=-1)
    # 获取batch size
    batch_size = tf.shape(y_true)[0]
    # 初始化总损失
    total_loss = tf.constant(0.0, dtype=tf.float32)
    # 对batch中的每个样本单独计算损失
    for i in tf.range(batch_size):
        # 获取当前样本
        current_true = y_true[i]
        current_pred = y_pred[i]
        # 计算方向图
        true_pattern, _, _ = phaseBit_2_pattern(current_true)
        pred_pattern, _, _ = phaseBit_2_pattern(current_pred)
        # 扩展维度以符合 SSIM 输入格式 (height, width) -> (1, height, width, 1)
        true_img = tf.expand_dims(tf.expand_dims(true_pattern, axis=0), axis=-1)
        pred_img = tf.expand_dims(tf.expand_dims(pred_pattern, axis=0), axis=-1)
        # SSIM
        ssim_val = tf.reduce_mean(tf.image.ssim(true_img, pred_img, max_val=1.0))
        loss = 1.0 - ssim_val  # 损失 = 1 - SSIM
        total_loss += loss
    avg_loss = total_loss / tf.cast(batch_size, tf.float32)
    return avg_loss
